{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a254e37b",
   "metadata": {},
   "source": [
    "## Trackio in Jupyter Notebooks\n",
    "\n",
    "This notebook demonstrates how to log training metrics with TrackIO, and display the results interactively directly in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fe95112",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q trackio"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daa46a91",
   "metadata": {},
   "source": [
    "The following code simulates training runs by generating synthetic loss, accuracy, and gradient norm metrics over multiple epochs and logs these metrics to Trackio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12cce6ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import random\n",
    "import time\n",
    "\n",
    "import trackio as wandb\n",
    "\n",
    "EPOCHS = 20\n",
    "PROJECT_ID = random.randint(100000, 999999)\n",
    "\n",
    "\n",
    "def generate_loss_curve(epoch, max_epochs, base_loss=2.5, min_loss=0.1):\n",
    "    \"\"\"Generate a realistic loss curve that decreases over time with noise\"\"\"\n",
    "    progress = epoch / max_epochs\n",
    "    base_curve = base_loss * math.exp(-3 * progress) + min_loss\n",
    "\n",
    "    noise_scale = 0.3 * (1 - progress * 0.7)\n",
    "    noise = random.gauss(0, noise_scale)\n",
    "\n",
    "    return max(min_loss * 0.5, base_curve + noise)\n",
    "\n",
    "\n",
    "def generate_accuracy_curve(epoch, max_epochs, max_acc=0.95, min_acc=0.1):\n",
    "    \"\"\"Generate a realistic accuracy curve that increases over time with noise\"\"\"\n",
    "    progress = epoch / max_epochs\n",
    "    base_curve = max_acc / (1 + math.exp(-6 * (progress - 0.5))) + min_acc\n",
    "\n",
    "    noise_scale = 0.08 * (1 - progress * 0.5)\n",
    "    noise = random.gauss(0, noise_scale)\n",
    "\n",
    "    return max(0, min(max_acc, base_curve + noise))\n",
    "\n",
    "\n",
    "def generate_grad_norm_curve(epoch, max_epochs):\n",
    "    \"\"\"Generate a gradient norm that starts at infinity and decreases to reasonable values\"\"\"\n",
    "    if epoch == 0:\n",
    "        return float(\"inf\")\n",
    "    elif epoch == 1:\n",
    "        return 1000.0\n",
    "    elif epoch == 2:\n",
    "        return 100.0\n",
    "    else:\n",
    "        progress = (epoch - 2) / (max_epochs - 2)\n",
    "        base_value = 50 * math.exp(-4 * progress) + 1.0\n",
    "        noise = random.gauss(0, 0.5)\n",
    "        return max(0.1, base_value + noise)\n",
    "\n",
    "\n",
    "for run in range(3):\n",
    "    wandb.init(\n",
    "        project=f\"fake-training-{PROJECT_ID}\",\n",
    "        name=f\"test-run-{run}\",\n",
    "        config=dict(\n",
    "            epochs=EPOCHS,\n",
    "            learning_rate=0.001,\n",
    "            batch_size=32,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    for epoch in range(EPOCHS):\n",
    "        train_loss = generate_loss_curve(\n",
    "            epoch,\n",
    "            EPOCHS,\n",
    "            base_loss=random.uniform(2.5, 3.5),\n",
    "            min_loss=random.uniform(0.05, 0.15),\n",
    "        )\n",
    "        val_loss = generate_loss_curve(\n",
    "            epoch,\n",
    "            EPOCHS,\n",
    "            base_loss=random.uniform(2.5, 3.5),\n",
    "            min_loss=random.uniform(0.05, 0.15),\n",
    "        )\n",
    "\n",
    "        train_accuracy = generate_accuracy_curve(\n",
    "            epoch,\n",
    "            EPOCHS,\n",
    "            max_acc=random.uniform(0.7, 0.9),\n",
    "            min_acc=random.uniform(0.1, 0.3),\n",
    "        )\n",
    "        val_accuracy = generate_accuracy_curve(\n",
    "            epoch,\n",
    "            EPOCHS,\n",
    "            max_acc=random.uniform(0.7, 0.9),\n",
    "            min_acc=random.uniform(0.1, 0.3),\n",
    "        )\n",
    "\n",
    "        grad_norm = generate_grad_norm_curve(epoch, EPOCHS)\n",
    "\n",
    "        if epoch > 2 and random.random() < 0.3:\n",
    "            val_loss *= 1.1\n",
    "            val_accuracy *= 0.95\n",
    "\n",
    "        wandb.log(\n",
    "            {\n",
    "                \"train/loss\": round(train_loss, 4),\n",
    "                \"train/accuracy\": round(train_accuracy, 4),\n",
    "                \"train/rewards/reward1\": random.random(),\n",
    "                \"train/rewards/reward2\": random.random(),\n",
    "                \"val/loss\": round(val_loss, 4),\n",
    "                \"val/accuracy\": round(val_accuracy, 4),\n",
    "                \"grad_norm\": grad_norm,\n",
    "            }\n",
    "        )\n",
    "\n",
    "        time.sleep(0.2)\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c471846d",
   "metadata": {},
   "source": [
    "The following cell launches the TrackIO dashboard directly in the notebook, allowing to interactively explore your logged training metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "957f6b6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import trackio\n",
    "\n",
    "trackio.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
